from collections.abc import Iterator
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset
from .bounded import BoundedDataset, BoundedIterator
from .delegating import DelegatingDataset, DelegatingIterator
from .skipped import SkippedDataset, SkippedIterator
from .striding import StridingDataset, StridingIterator

AnySample = TypeVar('AnySample')


class SlicedIterator(DelegatingIterator[AnySample]):
    def __init__(self, iterator: Iterator[AnySample], slice_: slice):
        start = slice_.start
        stop = slice_.stop
        step = slice_.step

        if start is not None:
            iterator = SkippedIterator(iterator, start)
        if step is not None:
            iterator = StridingIterator(iterator, step)
        if stop is not None:
            start = start if start is not None else 0
            step = step if step is not None else 1
            iterator = BoundedIterator(iterator, (stop - start + step - 1) // step)

        super().__init__(iterator)


class SlicedDataset(DelegatingDataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], slice_: slice):
        start = slice_.start
        stop = slice_.stop
        step = slice_.step

        if start is not None:
            dataset = SkippedDataset(dataset, start)
        if step is not None:
            dataset = StridingDataset(dataset, step)
        if stop is not None:
            start = start if start is not None else 0
            step = step if step is not None else 1
            dataset = BoundedDataset(dataset, (stop - start + step - 1) // step)

        super().__init__(dataset)
